from typing import Dict, Optional, Any

import torch
import torch_npu

import pytorch_lightning as pl
from pytorch_lightning.plugins.precision import PrecisionPlugin
from pytorch_lightning.strategies.single_device import SingleDeviceStrategy
from pytorch_lightning.utilities.exceptions import MisconfigurationException
from pytorch_lightning.utilities.types import _DEVICE
from pytorch_lightning.plugins.io.checkpoint_plugin import CheckpointIO

from lightning_npu.plugins.io.npu_plugin import NPUCheckpointIO


class SingleNPUStrategy(SingleDeviceStrategy):
    """Strategy for training on single NPU device."""

    strategy_name = "npu_single"

    def __init__(
        self,
        device: _DEVICE = "npu",
        accelerator: Optional["pl.accelerators.accelerator.Accelerator"] = None,
        precision_plugin: Optional[PrecisionPlugin] = None,
        device_index: int = None,
        checkpoint_io: Optional[CheckpointIO] = None,
    ):
        if device_index is not None:
            device = f"{device}:{device_index}"
        super().__init__(
            accelerator=accelerator,
            device=device,
            precision_plugin=precision_plugin,
            checkpoint_io=checkpoint_io or NPUCheckpointIO(),
        )

    @property
    def is_distributed(self) -> bool:
        return False

    def setup(self, trainer: "pl.Trainer") -> None:
        self.model_to_device()
        super().setup(trainer)

    def setup_optimizers(self, trainer: "pl.Trainer") -> None:
        super().setup_optimizers(trainer)

        if len(self.optimizers) > 1:
            raise MisconfigurationException("NPUs currently support only one optimizer.")

    def model_to_device(self) -> None:
        self.model.to(self.root_device)  # type: ignore

    @classmethod
    def register_strategies(cls, strategy_registry: Dict) -> None:
        strategy_registry.register(
            cls.strategy_name,
            cls,
            description=f"{cls.__class__.__name__}",
        )

    def load_checkpoint(self, checkpoint_path) -> Dict[str, Any]:
        torch.npu.empty_cache()
        current_device = torch.device("npu", torch.npu.current_device())
        return self.checkpoint_io.load_checkpoint(checkpoint_path, current_device)
    